{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dba3cf27",
   "metadata": {},
   "source": [
    "# Learn quantization with hands on"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eac2173",
   "metadata": {},
   "source": [
    "## Asymmetric Quantization\n",
    "非对称量化是一种神经网络量化方法，将浮点权重或激活值映射到整数范围时，允许零点（zero-point）偏移，即量化区间的最小值不一定对应浮点值中的0。这使得量化区间更灵活，能更好适应非对称分布的数据（如包含负数和正数），从而减少量化误差，提升模型精度。\n",
    "\n",
    "  ![alt text](resources/asym_quant.png \"Title\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ca246819",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import Tensor\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c78ecf9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_q_with_scale_and_zero_points(\n",
    "    tensor: Tensor, scale, zero_point, dtype=torch.int8\n",
    ") -> Tensor:\n",
    "    \"\"\"\n",
    "    Quantizes a tensor using the provided scale and zero point.\n",
    "\n",
    "    Args:\n",
    "        tensor (Tensor): The input tensor to quantize.\n",
    "        scale (float): The scale factor for quantization.\n",
    "        zero_point (int): The zero point for quantization.\n",
    "        dtype (torch.dtype, optional): The desired data type of the output tensor. Defaults to torch.int8.\n",
    "    \"\"\"\n",
    "    scaled_and_shifted_tensor = tensor / scale + zero_point\n",
    "    rounded_tensor = torch.round(scaled_and_shifted_tensor)\n",
    "\n",
    "    qmin = torch.iinfo(dtype).min\n",
    "    qmax = torch.iinfo(dtype).max\n",
    "\n",
    "    return rounded_tensor.clamp(qmin, qmax).to(dtype)\n",
    "\n",
    "\n",
    "def get_q_scale_and_zero_point(tensor: Tensor, dtype=torch.int8) -> (float, int):\n",
    "    \"\"\"\n",
    "    Computes the scale and zero point for quantization.\n",
    "\n",
    "    Args:\n",
    "        tensor (Tensor): The input tensor to compute scale and zero point.\n",
    "        dtype (torch.dtype, optional): The desired data type of the output tensor. Defaults to torch.int8.\n",
    "\n",
    "    Returns:\n",
    "        tuple: A tuple containing the scale and zero point.\n",
    "    \"\"\"\n",
    "    qmin = torch.iinfo(dtype).min\n",
    "    qmax = torch.iinfo(dtype).max\n",
    "    min_val = tensor.min().item()\n",
    "    max_val = tensor.max().item()\n",
    "    if min_val == max_val:\n",
    "        # If all values are the same, we can set scale to value and zero point to 0\n",
    "        return min_val, 0\n",
    "\n",
    "    scale = (max_val - min_val) / (qmax - qmin)\n",
    "    zero_point = int(round(qmin - min_val / scale))\n",
    "    zero_point = max(qmin, min(zero_point, qmax))\n",
    "\n",
    "    return scale, zero_point\n",
    "\n",
    "\n",
    "def linear_quantization(tensor: Tensor, dtype=torch.int8):\n",
    "    \"\"\"\n",
    "    Quantizes a tensor to the specified data type using linear quantization.\n",
    "\n",
    "    Args:\n",
    "        tensor (Tensor): The input tensor to quantize.\n",
    "        dtype (torch.dtype, optional): The desired data type of the output tensor. Defaults to torch.int8.\n",
    "\n",
    "    Returns:\n",
    "        Tensor: The quantized tensor.\n",
    "    \"\"\"\n",
    "    scale, zero_point = get_q_scale_and_zero_point(tensor, dtype)\n",
    "    return (\n",
    "        linear_q_with_scale_and_zero_points(tensor, scale, zero_point, dtype),\n",
    "        scale,\n",
    "        zero_point,\n",
    "    )\n",
    "\n",
    "\n",
    "def linear_dequantization(quantized_tensor: Tensor, scale, zero_point) -> Tensor:\n",
    "    return scale * (quantized_tensor.float() - zero_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c31b0e8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scale: 0.12882560281192554, Zero Point: 3\n",
      "Original Tensor:\n",
      " tensor([[ -6.4771,   0.7486,  -5.5430],\n",
      "        [-16.8492,  14.8348,  10.8760],\n",
      "        [  2.0805,   5.5973,  16.0013]])\n",
      "Quantized Tensor:\n",
      " tensor([[ -47,    9,  -40],\n",
      "        [-128,  118,   87],\n",
      "        [  19,   46,  127]], dtype=torch.int8)\n",
      "Dequantized Tensor:\n",
      " tensor([[ -6.4413,   0.7730,  -5.5395],\n",
      "        [-16.8762,  14.8149,  10.8214],\n",
      "        [  2.0612,   5.5395,  15.9744]])\n",
      "Error tensor([[ 0.0359,  0.0243,  0.0035],\n",
      "        [-0.0269, -0.0199, -0.0546],\n",
      "        [-0.0193, -0.0578, -0.0269]])\n",
      "Relative error:\n",
      " tensor([[-0.0055,  0.0325, -0.0006],\n",
      "        [ 0.0016, -0.0013, -0.0050],\n",
      "        [-0.0093, -0.0103, -0.0017]])\n",
      "MSE:\n",
      " tensor(0.0012)\n"
     ]
    }
   ],
   "source": [
    "### a dummy tensor to test the implementation\n",
    "test_tensor = torch.randn(3, 3) * 10\n",
    "quatization_type = torch.int8\n",
    "quantized_tensor, scale, zero_point = linear_quantization(\n",
    "    test_tensor, dtype=quatization_type\n",
    ")\n",
    "print(f\"Scale: {scale}, Zero Point: {zero_point}\")\n",
    "print(\"Original Tensor:\\n\", test_tensor)\n",
    "print(\"Quantized Tensor:\\n\", quantized_tensor)\n",
    "dequantized_tensor = linear_dequantization(quantized_tensor, scale, zero_point)\n",
    "print(\"Dequantized Tensor:\\n\", dequantized_tensor)\n",
    "print(\"Error\", dequantized_tensor - test_tensor)\n",
    "print(\"Relative error:\\n\", (dequantized_tensor - test_tensor) / (test_tensor + 1e-9))\n",
    "print(\"MSE:\\n\", (dequantized_tensor - test_tensor).square().mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7c41626",
   "metadata": {},
   "source": [
    "## Symmetric Quantization\n",
    "对称量化是一种神经网络量化方法，将浮点值对称地映射到整数范围，要求量化区间的最小值和最大值关于零对称（如[-a, a]），且浮点零映射到整数零点。它不引入零点偏移，计算更简单高效，但适用于数据分布接近对称的情况，常用于权重量化。\n",
    "\n",
    "  ![alt text](resources/sym_quant.png \"Title\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "892acd42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_q_scale_symmetric(tensor: Tensor, dtype=torch.int8) -> float:\n",
    "    \"\"\"\n",
    "    Computes the symmetric scale for quantization.\n",
    "\n",
    "    Args:\n",
    "        tensor (Tensor): The input tensor to compute scale.\n",
    "        dtype (torch.dtype, optional): The desired data type of the output tensor. Defaults to torch.int8.\n",
    "\n",
    "    Returns:\n",
    "        float: The symmetric scale.\n",
    "    \"\"\"\n",
    "    tensor_max = tensor.abs().max().item()\n",
    "    qmax = torch.iinfo(dtype).max\n",
    "    return tensor_max / qmax if tensor_max != 0 else 1.0\n",
    "\n",
    "\n",
    "def linear_q_symmetric(tensor: Tensor, dtype=torch.int8):\n",
    "    \"\"\"\n",
    "    Quantizes a tensor using symmetric quantization.\n",
    "\n",
    "    Args:\n",
    "        tensor (Tensor): The input tensor to quantize.\n",
    "        dtype (torch.dtype, optional): The desired data type of the output tensor. Defaults to torch.int8.\n",
    "\n",
    "    Returns:\n",
    "        Tensor: The quantized tensor.\n",
    "    \"\"\"\n",
    "    scale = get_q_scale_symmetric(tensor, dtype)\n",
    "    quantized_tensor = linear_q_with_scale_and_zero_points(\n",
    "        tensor, scale=scale, zero_point=0, dtype=dtype\n",
    "    )\n",
    "    return quantized_tensor, scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "429d22b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Symmetric Scale: 0.17191350741649236\n",
      "Original Tensor:\n",
      " tensor([[  4.3260,  -6.3765,  -5.5382,  -4.1744],\n",
      "        [  3.3907,   2.1706,   3.6630,  -9.7095],\n",
      "        [ 21.2197,  -3.7216, -14.5955, -12.1297],\n",
      "        [-21.8330,  -3.5494, -14.7311, -14.0608]])\n",
      "Quantized Tensor:\n",
      " tensor([[  25,  -37,  -32,  -24],\n",
      "        [  20,   13,   21,  -56],\n",
      "        [ 123,  -22,  -85,  -71],\n",
      "        [-127,  -21,  -86,  -82]], dtype=torch.int8)\n",
      "Dequantized Tensor:\n",
      " tensor([[  4.2978,  -6.3608,  -5.5012,  -4.1259],\n",
      "        [  3.4383,   2.2349,   3.6102,  -9.6272],\n",
      "        [ 21.1454,  -3.7821, -14.6126, -12.2059],\n",
      "        [-21.8330,  -3.6102, -14.7846, -14.0969]])\n",
      "Error tensor([[-0.0281,  0.0157,  0.0369,  0.0485],\n",
      "        [ 0.0476,  0.0643, -0.0528,  0.0824],\n",
      "        [-0.0744, -0.0605, -0.0171, -0.0762],\n",
      "        [ 0.0000, -0.0608, -0.0535, -0.0361]])\n",
      "Relative error:\n",
      " tensor([[-0.0065, -0.0025, -0.0067, -0.0116],\n",
      "        [ 0.0140,  0.0296, -0.0144, -0.0085],\n",
      "        [-0.0035,  0.0163,  0.0012,  0.0063],\n",
      "        [-0.0000,  0.0171,  0.0036,  0.0026]])\n",
      "MSE:\n",
      " tensor(0.0027)\n",
      "Max absolute error: tensor(0.0824)\n"
     ]
    }
   ],
   "source": [
    "test_tensor = torch.randn(4, 4) * 10\n",
    "quantized_tensor, scale = linear_q_symmetric(test_tensor, dtype=torch.int8)\n",
    "print(f\"Symmetric Scale: {scale}\")\n",
    "print(\"Original Tensor:\\n\", test_tensor)\n",
    "print(\"Quantized Tensor:\\n\", quantized_tensor)\n",
    "dequantized_tensor = linear_dequantization(quantized_tensor, scale, 0)\n",
    "print(\"Dequantized Tensor:\\n\", dequantized_tensor)\n",
    "print(\"Error\", dequantized_tensor - test_tensor)\n",
    "print(\"Relative error:\\n\", (dequantized_tensor - test_tensor) / (test_tensor + 1e-9))\n",
    "print(\"MSE:\\n\", (dequantized_tensor - test_tensor).square().mean())\n",
    "print(\"Max absolute error:\", (dequantized_tensor - test_tensor).abs().max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c8aa46b",
   "metadata": {},
   "source": [
    "## Per Channel Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9a0a655e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tf/anaconda3/envs/pytorch-cu124/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from timm.layers import linear\n",
    "\n",
    "\n",
    "def linear_q_symmetric_per_channel(tensor: Tensor, dim: int, dtype=torch.int8):\n",
    "    \"\"\"Do symmetric quantization per channel.\n",
    "\n",
    "    Args:\n",
    "        tensor: input Tensor.\n",
    "        dim: Dim to quantize over.\n",
    "        dtype: _description_. Defaults to torch.int8.\n",
    "    \"\"\"\n",
    "    output_dim = tensor.shape[dim]\n",
    "    scales = torch.zeros(output_dim, dtype=torch.float32)\n",
    "    for index in range(output_dim):\n",
    "        sub_tensor = tensor.select(dim, index)\n",
    "        scales[index] = get_q_scale_symmetric(sub_tensor, dtype)\n",
    "    scale_shape = [1] * tensor.dim()\n",
    "    scale_shape[dim] = output_dim\n",
    "    scales = scales.view(scale_shape)\n",
    "    quantized_tensor = linear_q_with_scale_and_zero_points(\n",
    "        tensor, scales, zero_point=0, dtype=dtype\n",
    "    )\n",
    "    return quantized_tensor, scales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d47de1ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Symmetric Per Channel Scale: tensor([[5.7370],\n",
      "        [2.3268],\n",
      "        [5.3906]])\n",
      "Original Tensor:\n",
      " tensor([[ 191.6000,  -13.5000,  728.6000],\n",
      "        [  92.1400,  295.5000, -184.0000],\n",
      "        [   0.0000,  684.6000,  245.5000]])\n",
      "Quantized Tensor:\n",
      " tensor([[ 33,  -2, 127],\n",
      "        [ 40, 127, -79],\n",
      "        [  0, 127,  46]], dtype=torch.int8)\n",
      "Dequantized Tensor:\n",
      " tensor([[ 189.3213,  -11.4740,  728.6000],\n",
      "        [  93.0709,  295.5000, -183.8150],\n",
      "        [   0.0000,  684.6000,  247.9653]])\n",
      "Error tensor([[-2.2787,  2.0260,  0.0000],\n",
      "        [ 0.9309,  0.0000,  0.1850],\n",
      "        [ 0.0000,  0.0000,  2.4653]])\n",
      "Relative error:\n",
      " tensor([[-0.0119, -0.1501,  0.0000],\n",
      "        [ 0.0101,  0.0000, -0.0010],\n",
      "        [ 0.0000,  0.0000,  0.0100]])\n",
      "MSE:\n",
      " tensor(1.8084)\n",
      "Max absolute error: tensor(2.4653)\n"
     ]
    }
   ],
   "source": [
    "test_tensor = torch.tensor(\n",
    "    [\n",
    "        [191.6, -13.5, 728.6],\n",
    "        [92.14, 295.5, -184],\n",
    "        [0, 684.6, 245.5],\n",
    "    ]\n",
    ")\n",
    "quantized_tensor, scale = linear_q_symmetric_per_channel(\n",
    "    test_tensor, dim=0, dtype=torch.int8\n",
    ")\n",
    "print(f\"Symmetric Per Channel Scale: {scale}\")\n",
    "print(\"Original Tensor:\\n\", test_tensor)\n",
    "print(\"Quantized Tensor:\\n\", quantized_tensor)\n",
    "dequantized_tensor = linear_dequantization(quantized_tensor, scale, 0)\n",
    "print(\"Dequantized Tensor:\\n\", dequantized_tensor)\n",
    "print(\"Error\", dequantized_tensor - test_tensor)\n",
    "print(\"Relative error:\\n\", (dequantized_tensor - test_tensor) / (test_tensor + 1e-9))\n",
    "print(\"MSE:\\n\", (dequantized_tensor - test_tensor).square().mean())\n",
    "print(\"Max absolute error:\", (dequantized_tensor - test_tensor).abs().max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "79aad30a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Symmetric Per Channel Scale: tensor([[1.5087, 5.3906, 5.7370]])\n",
      "Original Tensor:\n",
      " tensor([[ 191.6000,  -13.5000,  728.6000],\n",
      "        [  92.1400,  295.5000, -184.0000],\n",
      "        [   0.0000,  684.6000,  245.5000]])\n",
      "Quantized Tensor:\n",
      " tensor([[127,  -3, 127],\n",
      "        [ 61,  55, -32],\n",
      "        [  0, 127,  43]], dtype=torch.int8)\n",
      "Dequantized Tensor:\n",
      " tensor([[ 191.6000,  -16.1717,  728.6000],\n",
      "        [  92.0284,  296.4803, -183.5842],\n",
      "        [   0.0000,  684.6000,  246.6913]])\n",
      "Error tensor([[ 0.0000, -2.6717,  0.0000],\n",
      "        [-0.1116,  0.9803,  0.4158],\n",
      "        [ 0.0000,  0.0000,  1.1913]])\n",
      "Relative error:\n",
      " tensor([[ 0.0000,  0.1979,  0.0000],\n",
      "        [-0.0012,  0.0033, -0.0023],\n",
      "        [ 0.0000,  0.0000,  0.0049]])\n",
      "MSE:\n",
      " tensor(1.0781)\n",
      "Max absolute error: tensor(2.6717)\n"
     ]
    }
   ],
   "source": [
    "quantized_tensor, scale = linear_q_symmetric_per_channel(\n",
    "    test_tensor, dim=1, dtype=torch.int8\n",
    ")\n",
    "print(f\"Symmetric Per Channel Scale: {scale}\")\n",
    "print(\"Original Tensor:\\n\", test_tensor)\n",
    "print(\"Quantized Tensor:\\n\", quantized_tensor)\n",
    "dequantized_tensor = linear_dequantization(quantized_tensor, scale, 0)\n",
    "print(\"Dequantized Tensor:\\n\", dequantized_tensor)\n",
    "print(\"Error\", dequantized_tensor - test_tensor)\n",
    "print(\"Relative error:\\n\", (dequantized_tensor - test_tensor) / (test_tensor + 1e-9))\n",
    "print(\"MSE:\\n\", (dequantized_tensor - test_tensor).square().mean())\n",
    "print(\"Max absolute error:\", (dequantized_tensor - test_tensor).abs().max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f1bb8f1",
   "metadata": {},
   "source": [
    "## Per Group Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cf4e9342",
   "metadata": {},
   "outputs": [],
   "source": [
    "from click import group\n",
    "\n",
    "\n",
    "def linear_q_symmetric_per_group(tensor: Tensor, group_size: int, dtype=torch.int8):\n",
    "    \"\"\"Do symmetric quantization per group.\n",
    "\n",
    "    Args:\n",
    "        tensor: input Tensor.\n",
    "        group_size: Size of the group to quantize over.\n",
    "        dtype: _description_. Defaults to torch.int8.\n",
    "    \"\"\"\n",
    "    t_shape = tensor.shape\n",
    "    assert t_shape[1] % group_size == 0, \"Tensor size must be divisible by group_size\"\n",
    "    assert tensor.dim() == 2, \"Tensor must be 2D for group quantization\"\n",
    "\n",
    "    tensor = tensor.view(-1, group_size)\n",
    "    quantized_tensor, scale = linear_q_symmetric_per_channel(tensor, dim=0, dtype=dtype)\n",
    "    quantized_tensor = quantized_tensor.view(t_shape)\n",
    "\n",
    "    return quantized_tensor, scale\n",
    "\n",
    "\n",
    "def linear_dequantization_per_group(\n",
    "    quantized_tensor: Tensor, scale, group_size: int\n",
    ") -> Tensor:\n",
    "    \"\"\"Dequantizes a tensor that was quantized per group.\n",
    "\n",
    "    Args:\n",
    "        quantized_tensor: The quantized tensor to dequantize.\n",
    "        scale: The scale used for quantization.\n",
    "        group_size: Size of the group used for quantization.\n",
    "    \"\"\"\n",
    "    t_shape = quantized_tensor.shape\n",
    "    assert t_shape[1] % group_size == 0, \"Tensor size must be divisible by group_size\"\n",
    "    assert quantized_tensor.dim() == 2, \"Tensor must be 2D for group dequantization\"\n",
    "    quantized_tensor = quantized_tensor.view(-1, group_size)\n",
    "    dequantized_tensor = linear_dequantization(quantized_tensor, scale, 0)\n",
    "    dequantized_tensor = dequantized_tensor.view(t_shape)\n",
    "\n",
    "    return dequantized_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c33e1d94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Symmetric Per Group Scale: tensor([[0.1412],\n",
      "        [0.1113],\n",
      "        [0.0864],\n",
      "        [0.0295],\n",
      "        [0.1109],\n",
      "        [0.1820],\n",
      "        [0.1905],\n",
      "        [0.0663],\n",
      "        [0.1533],\n",
      "        [0.0806],\n",
      "        [0.0367],\n",
      "        [0.1105]])\n",
      "Original Tensor:\n",
      " tensor([[ 17.4175,  17.3003, -17.9289, -10.8549,  -3.3364,  14.1369],\n",
      "        [ 10.9709,   8.9511,   3.1999,   3.7452,   3.1249,  -0.2184],\n",
      "        [-12.7733,  10.0035,  14.0848,  -3.9805, -23.1181,  11.8203],\n",
      "        [  1.1131, -24.1891,  -9.8270,  -7.5528,   8.4201,  -6.2769],\n",
      "        [-19.4720, -11.6513,  -3.5473, -10.2308,  -2.6917,  -0.8038],\n",
      "        [  0.3371,   4.6565,   1.7423,  -8.6143,   5.0310, -14.0398]])\n",
      "Quantized Tensor:\n",
      " tensor([[ 123,  123, -127,  -98,  -30,  127],\n",
      "        [ 127,  104,   37,  127,  106,   -7],\n",
      "        [-115,   90,  127,  -22, -127,   65],\n",
      "        [   6, -127,  -52, -114,  127,  -95],\n",
      "        [-127,  -76,  -23, -127,  -33,  -10],\n",
      "        [   9,  127,   48,  -78,   46, -127]], dtype=torch.int8)\n",
      "Dequantized Tensor:\n",
      " tensor([[ 17.3642,  17.3642, -17.9289, -10.9088,  -3.3394,  14.1369],\n",
      "        [ 10.9709,   8.9841,   3.1963,   3.7452,   3.1259,  -0.2064],\n",
      "        [-12.7540,   9.9814,  14.0848,  -4.0047, -23.1181,  11.8321],\n",
      "        [  1.1428, -24.1891,  -9.9042,  -7.5582,   8.4201,  -6.2985],\n",
      "        [-19.4720, -11.6525,  -3.5264, -10.2308,  -2.6584,  -0.8056],\n",
      "        [  0.3300,   4.6565,   1.7599,  -8.6229,   5.0853, -14.0398]])\n",
      "Error tensor([[-0.0533,  0.0639,  0.0000, -0.0539, -0.0030,  0.0000],\n",
      "        [ 0.0000,  0.0329, -0.0037,  0.0000,  0.0010,  0.0120],\n",
      "        [ 0.0193, -0.0221,  0.0000, -0.0242,  0.0000,  0.0118],\n",
      "        [ 0.0297,  0.0000, -0.0772, -0.0055,  0.0000, -0.0216],\n",
      "        [ 0.0000, -0.0013,  0.0209,  0.0000,  0.0333, -0.0018],\n",
      "        [-0.0072,  0.0000,  0.0176, -0.0086,  0.0543,  0.0000]])\n",
      "Relative error:\n",
      " tensor([[-0.0031,  0.0037, -0.0000,  0.0050,  0.0009,  0.0000],\n",
      "        [ 0.0000,  0.0037, -0.0011,  0.0000,  0.0003, -0.0550],\n",
      "        [-0.0015, -0.0022,  0.0000,  0.0061, -0.0000,  0.0010],\n",
      "        [ 0.0267, -0.0000,  0.0079,  0.0007,  0.0000,  0.0034],\n",
      "        [-0.0000,  0.0001, -0.0059, -0.0000, -0.0124,  0.0023],\n",
      "        [-0.0212,  0.0000,  0.0101,  0.0010,  0.0108, -0.0000]])\n",
      "MSE:\n",
      " tensor(0.0007)\n",
      "Max absolute error: tensor(0.0772)\n"
     ]
    }
   ],
   "source": [
    "test_tensor = torch.randn(6, 6) * 10\n",
    "group_size = 3\n",
    "quantized_tensor, scale = linear_q_symmetric_per_group(\n",
    "    test_tensor, group_size=group_size, dtype=torch.int8\n",
    ")\n",
    "print(f\"Symmetric Per Group Scale: {scale}\")\n",
    "print(\"Original Tensor:\\n\", test_tensor)\n",
    "print(\"Quantized Tensor:\\n\", quantized_tensor)\n",
    "dequantized_tensor = linear_dequantization_per_group(\n",
    "    quantized_tensor, scale, group_size\n",
    ")\n",
    "print(\"Dequantized Tensor:\\n\", dequantized_tensor)\n",
    "print(\"Error\", dequantized_tensor - test_tensor)\n",
    "print(\"Relative error:\\n\", (dequantized_tensor - test_tensor) / (test_tensor + 1e-9))\n",
    "print(\"MSE:\\n\", (dequantized_tensor - test_tensor).square().mean())\n",
    "print(\"Max absolute error:\", (dequantized_tensor - test_tensor).abs().max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b83ee69",
   "metadata": {},
   "source": [
    "## Inference Linear Quantization W8A32\n",
    "\n",
    "Weight only quantization and activation is performed with float32."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "81285859",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output with quantized weights:\n",
      " tensor([-2.9965,  3.8768,  9.3957])\n",
      "Output with original weights:\n",
      " tensor([-3.0000,  3.8500,  9.3800])\n",
      "Difference:\n",
      " tensor([0.0035, 0.0268, 0.0157])\n",
      "Relative error:\n",
      " tensor([-0.0012,  0.0070,  0.0017])\n",
      "MSE:\n",
      " tensor(0.0003)\n",
      "Max absolute error: tensor(0.0268)\n"
     ]
    }
   ],
   "source": [
    "def quantized_linear_W8A32_without_bias(input: Tensor, q_w, s_w, z_w) -> Tensor:\n",
    "    \"\"\"\n",
    "    Computes the output of a linear layer with weight quantization.\n",
    "\n",
    "    Args:\n",
    "        input (Tensor): The input tensor.\n",
    "        q_w (Tensor): The quantized weight tensor.\n",
    "        s_w (float): The scale for the weights.\n",
    "        z_w (int): The zero point for the weights.\n",
    "\n",
    "    Returns:\n",
    "        Tensor: The output tensor after applying the quantized linear transformation.\n",
    "    \"\"\"\n",
    "    assert (\n",
    "        input.dtype == torch.float32\n",
    "    ), \"Input tensor must be float32 for linear operation\"\n",
    "    assert q_w.dtype == torch.int8, \"Quantized weights must be int8\"\n",
    "\n",
    "    # Dequantize weights\n",
    "    w = linear_dequantization(q_w, s_w, z_w)\n",
    "    # Perform linear operation\n",
    "    return torch.nn.functional.linear(input, w, bias=None)\n",
    "\n",
    "\n",
    "input = torch.tensor([1, 2, 3], dtype=torch.float32)\n",
    "weight = torch.tensor([[-2, -1.13, 0.42], [-1.51, 0.25, 1.62], [0.23, 1.35, 2.15]])\n",
    "q_w, s_w = linear_q_symmetric(weight)\n",
    "output_quant = quantized_linear_W8A32_without_bias(input, q_w, s_w, 0)\n",
    "output_origin = torch.nn.functional.linear(input, weight, bias=None)\n",
    "print(\"Output with quantized weights:\\n\", output_quant)\n",
    "print(\"Output with original weights:\\n\", output_origin)\n",
    "print(\"Difference:\\n\", output_quant - output_origin)\n",
    "print(\"Relative error:\\n\", (output_quant - output_origin) / (output_origin + 1e-9))\n",
    "print(\"MSE:\\n\", (output_quant - output_origin).square().mean())\n",
    "print(\"Max absolute error:\", (output_quant - output_origin).abs().max())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-cu124",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
